import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score as sklearn_accuracy_score  # 修正导入
import matplotlib.pyplot as plt
import os
import time
from typing import Tuple, List, Dict, Union

# 类型注解：明确函数输入输出类型
def load_data() -> Tuple[np.ndarray, np.ndarray]:
    """加载数据集，优先使用自定义路径数据，否则使用内置鸢尾花数据集"""
    try:
        data_path: str = r"D:\exp5\irisdata.csv"  # 修正路径为exp5
        if os.path.exists(data_path):
            df: pd.DataFrame = pd.read_csv(data_path)
            # 校验数据格式：至少包含2列（特征+标签）
            if df.shape[1] < 2:
                raise ValueError("数据集列数不足，需包含特征和标签")
            X: np.ndarray = df.iloc[:, :-1].values.astype(np.float64)
            y: np.ndarray = df.iloc[:, -1].values
            print(f"成功加载自定义数据集，形状: {X.shape}")
            return X, y
        print("未找到自定义数据集，将使用sklearn内置的iris数据集")
        iris = load_iris()
        return iris.data.astype(np.float64), iris.target
    except Exception as e:
        print(f"数据加载错误: {str(e)}，使用内置数据集")
        iris = load_iris()
        return iris.data.astype(np.float64), iris.target

# 自定义KNN分类器，增加参数校验
class KNNClassifier:
    def __init__(self, k: int = 5, distance_metric: str = 'euclidean'):
        # 校验k值合理性
        if not isinstance(k, int) or k < 1:
            raise ValueError("k必须为正整数")
        # 校验距离度量合法性
        valid_metrics: List[str] = ['euclidean', 'manhattan', 'chebyshev']
        if distance_metric not in valid_metrics:
            raise ValueError(f"不支持的距离度量，允许值: {valid_metrics}")
            
        self.k: int = k
        self.distance_metric: str = distance_metric
        self.X_train: Optional[np.ndarray] = None
        self.y_train: Optional[np.ndarray] = None

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """拟合模型，校验数据维度"""
        if X.shape[0] != len(y):
            raise ValueError("特征与标签样本数不一致")
        self.X_train = X
        self.y_train = y

    def predict(self, X: np.ndarray) -> np.ndarray:
        """批量预测，支持二维数组输入"""
        if self.X_train is None or self.y_train is None:
            raise RuntimeError("模型未训练，请先调用fit方法")
        return np.array([self._predict(x) for x in X])

    def _predict(self, x: np.ndarray) -> int:
        """单样本预测，核心逻辑优化"""
        if self.distance_metric == 'euclidean':
            distances: np.ndarray = np.linalg.norm(self.X_train - x, axis=1)
        elif self.distance_metric == 'manhattan':
            distances: np.ndarray = np.sum(np.abs(self.X_train - x), axis=1)
        else:  # chebyshev
            distances: np.ndarray = np.max(np.abs(self.X_train - x), axis=1)
            
        # 使用argsort获取前k个最小距离索引
        k_indices: np.ndarray = np.argsort(distances)[:self.k]
        # 统计标签频率，处理平局情况（取最小标签）
        labels: np.ndarray = self.y_train[k_indices]
        unique_labels, counts = np.unique(labels, return_counts=True)
        max_count: int = np.max(counts)
        # 多个标签同票时选择值最小的
        most_common = unique_labels[np.where(counts == max_count)[0][0]]
        return most_common

# 增加日志和可视化配置参数
def find_best_k(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
    max_k: int = 20,
    distance_metric: str = 'euclidean',
    log_file: Union[str, None] = None,
    save_figure: bool = True
) -> Tuple[int, float]:
    """寻找最优K值，支持日志记录和图像保存"""
    if max_k < 1:
        raise ValueError("max_k必须至少为1")
    accuracies: List[float] = []
    best_k: int = 1
    best_accuracy: float = 0.0

    for k in range(1, max_k + 1):
        knn = KNNClassifier(k=k, distance_metric=distance_metric)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        acc: float = sklearn_accuracy_score(y_test, y_pred)  # 修正函数名
        accuracies.append(acc)
        if acc > best_accuracy:
            best_accuracy = acc
            best_k = k

    # 记录详细日志
    if log_file:
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"\n=== {distance_metric.upper()}距离 K值搜索结果 ===")
            f.write(f"\nK值范围: 1-{max_k}, 测试集大小: {len(X_test)}")
            f.write("\nK值\t准确率")
            for k, acc in enumerate(accuracies, 1):
                f.write(f"\n{k}\t{acc:.4f}")
            f.write(f"\n最优K值: {best_k}, 最高准确率: {best_accuracy:.4f}\n")

    # 可视化配置优化
    plt.figure(figsize=(10, 6), dpi=300)
    plt.plot(range(1, max_k+1), accuracies, 'o-', linewidth=2, markersize=6)
    plt.title(f'K值-准确率曲线 ({distance_metric}距离)', fontsize=14)
    plt.xlabel('K值', fontsize=12)
    plt.ylabel('分类准确率', fontsize=12)
    plt.xticks(range(1, max_k+1), fontsize=10)
    plt.yticks(np.arange(0.8, 1.01, 0.02), fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    
    if save_figure and log_file:
        fig_path = log_file.replace('.log', f'_{distance_metric}_accuracy.png')
        plt.savefig(fig_path, bbox_inches='tight')
    # plt.show()  # 生产环境建议关闭图像显示

    return best_k, best_accuracy

def main():
    """主函数，包含完整实验流程控制"""
    # 生成带时间戳的日志文件（修正路径为exp5）
    timestamp: str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
    log_file: str = f"D:/exp5/knn_experiment_{timestamp}.log"  # 修正路径
    
    # 初始化日志
    with open(log_file, 'w', encoding='utf-8') as f:
        f.write(f"鸢尾花KNN分类实验日志 - {timestamp}\n")
        f.write(f"{'='*60}\n")
        f.write(f"实验配置:\n")
        f.write(f"数据集路径: D:\\exp5\\irisdata.csv\n")  # 修正路径
        f.write(f"测试集比例: 30%\n")
        f.write(f"距离度量: 欧式/曼哈顿/切比雪夫\n")
        f.write(f"{'='*60}\n\n")

    # 加载并校验数据
    X, y = load_data()
    with open(log_file, 'a', encoding='utf-8') as f:
        f.write(f"数据集详情:\n")
        f.write(f"样本总数: {len(X)}, 特征数: {X.shape[1]}\n")
        f.write(f"类别分布: {dict(zip(*np.unique(y, return_counts=True)))}\n\n")

    # 划分数据集（固定随机种子确保可复现）
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42, stratify=y
    )
    with open(log_file, 'a', encoding='utf-8') as f:
        f.write(f"数据划分结果:\n")
        f.write(f"训练集: {len(X_train)}样本, 测试集: {len(X_test)}样本\n\n")

    # 多距离度量实验
    distance_metrics: List[str] = ['euclidean', 'manhattan', 'chebyshev']
    results: Dict[str, Dict[str, Union[int, float]]] = {}

    for metric in distance_metrics:
        print(f"\n=== 正在执行 {metric.upper()} 距离实验 ===")
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"{'='*20} {metric.upper()}距离实验 {'='*20}\n")
        
        best_k, accuracy = find_best_k(
            X_train, y_train, X_test, y_test,
            max_k=20,
            distance_metric=metric,
            log_file=log_file,
            save_figure=True
        )
        results[metric] = {'best_k': best_k, 'accuracy': accuracy}

        # 模型对比实验
        # 自定义模型预测
        custom_knn = KNNClassifier(k=best_k, distance_metric=metric)
        custom_knn.fit(X_train, y_train)
        custom_pred = custom_knn.predict(X_test)
        
        # sklearn模型预测
        sklearn_knn = KNeighborsClassifier(
            n_neighbors=best_k,
            metric=metric,
            weights='uniform',
            algorithm='auto'
        )
        sklearn_knn.fit(X_train, y_train)
        sklearn_pred = sklearn_knn.predict(X_test)
        
        # 结果对比（使用正确的函数名）
        sklearn_accuracy = sklearn_accuracy_score(y_test, sklearn_pred)
        acc_diff: float = np.abs(accuracy - sklearn_accuracy)
        result_match: bool = np.array_equal(custom_pred, sklearn_pred)
        
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"\n模型对比结果 ({metric}距离):\n")
            f.write(f"自定义KNN: k={best_k}, 准确率={accuracy:.4f}\n")
            f.write(f"sklearn KNN: 准确率={sklearn_accuracy:.4f}\n")
            f.write(f"准确率差异: {acc_diff:.6f}\n")
            f.write(f"预测结果一致性: {result_match}\n\n")

    # 实验总结
    with open(log_file, 'a', encoding='utf-8') as f:
        f.write(f"{'='*60}\n")
        f.write(f"实验结论:\n")
        for metric, result in results.items():
            f.write(f"{metric.upper()}距离 - 最优k={result['best_k']}, 准确率={result['accuracy']:.4f}\n")
        f.write(f"\n完整实验日志路径: {log_file}\n")
        f.write(f"{'='*60}\n")

    print(f"\n实验完成！结果已保存至 {log_file}")
    print("各距离度量最优K值及准确率:")
    for metric, result in results.items():
        print(f"{metric.upper()}: k={result['best_k']}, 准确率={result['accuracy']:.4f}")

if __name__ == "__main__":
    main()